#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pandas as pd
import matplotlib.pyplot as plt
import re
import argparse
from datetime import datetime
import sys

def parse_topdump(file_path):
    """
    Parses the topdump.txt file to extract CPU usage information for the main process and its threads.

    Parameters:
    - file_path (str): Path to the topdump.txt file.

    Returns:
    - pandas.DataFrame: DataFrame containing relative time, PID, user, CPU usage, and command name.
    - dict: Dictionary mapping PID to command name.
    - str: Process name.
    - int: PID of the main process.
    """
    records = []
    pid_to_cmd = {}
    process_name = ""
    main_pid = None

    try:
        with open(file_path, 'r', encoding='utf-8', errors='replace') as file:
            lines = file.readlines()
    except UnicodeDecodeError:
        # If utf-8 fails, try using gbk encoding
        try:
            with open(file_path, 'r', encoding='gbk', errors='replace') as file:
                lines = file.readlines()
        except Exception as e:
            raise Exception(f"Failed to read file: {e}")

    current_section = None  # Current section being parsed: 'ps', 'top', or 'Process_Info'
    current_ps_pid = None
    spids = set()
    base_time = None
    relative_time = 0.0
    in_process_list = False  # Flag to indicate if we're in the process list of 'top'

    # Flags to identify if Process_Info has been parsed
    process_info_parsed = False

    for line_number, line in enumerate(lines, 1):
        line = line.strip()

        # Skip empty lines
        if not line:
            continue

        # Detect the start of the ps -T -p <PID> section
        ps_start_match = re.match(r'^===\s*ps\s+-T\s+-p\s+(\d+)\s+===', line)
        if ps_start_match:
            current_section = 'ps'
            current_ps_pid = int(ps_start_match.group(1))
            spids.clear()  # Reset SPIDs for new ps section
            continue

        # Detect the start of the Process_Info section
        process_info_start_match = re.match(r'^===\s*Process_Info\s*===', line)
        if process_info_start_match:
            current_section = 'Process_Info'
            continue

        # Detect the start of the top section
        top_start_match = re.match(r'^===\s*top\s+-H\s+-b\s+-n\s+\d+\s+-d\s+\d+\s+-p\s+\d+\s*===', line)
        if top_start_match:
            current_section = 'top'
            in_process_list = False  # Reset process list flag
            continue

        # Detect the end of any section
        if line.startswith('===') and not (ps_start_match or process_info_start_match or top_start_match):
            current_section = None
            continue

        # Parse the Process_Info section to get main process PID and name
        if current_section == 'Process_Info':
            # Skip headers and separator lines if any
            if re.match(r'^PID\s+TTY\s+TIME\s+CMD$', line) or re.match(r'^-+$', line):
                continue

            # Parse the process info line
            proc_info = re.split(r'\s+', line, maxsplit=3)
            if len(proc_info) < 4:
                continue
            try:
                pid = int(proc_info[0])
                # tty = proc_info[1]
                # time = proc_info[2]
                cmd = proc_info[3]
            except ValueError:
                continue

            main_pid = pid
            process_name = cmd
            process_info_parsed = True
            current_section = None  # Assuming only one line in Process_Info
            continue

        # Parse the ps section to extract SPIDs and map PIDs to command names
        if current_section == 'ps':
            # Skip headers and separator lines
            if re.match(r'^PID\s+SPID\s+TTY\s+TIME\s+CMD$', line) or re.match(r'^-+$', line):
                continue

            # Parse process and thread lines
            proc_info = re.split(r'\s+', line, maxsplit=4)
            if len(proc_info) < 5:
                continue
            try:
                pid = int(proc_info[0])
                spid = int(proc_info[1])
                cmd = proc_info[4]
            except ValueError:
                continue

            # Only consider the ps section corresponding to the main_pid
            if main_pid and current_ps_pid != main_pid:
                continue  # Skip irrelevant ps sections

            # Add SPID to the set
            spids.add(spid)
            # Map SPID to command name
            pid_to_cmd[spid] = cmd

            # If the SPID matches main_pid, set process_name (redundant if already set via Process_Info)
            if spid == main_pid and not process_info_parsed:
                process_name = cmd

            continue

        # Detect the start of the top -H -b -n <n> -d <d> -p <PID> section
        if line.startswith('top -'):
            current_section = 'top'
            in_process_list = False  # Reset process list flag
            # Extract timestamp
            timestamp_match = re.search(r'top\s+-\s+(\d{2}:\d{2}:\d{2})', line)
            if timestamp_match:
                timestamp_str = timestamp_match.group(1)
                try:
                    timestamp_dt = datetime.strptime(timestamp_str, '%H:%M:%S')
                    if base_time is None:
                        base_time = timestamp_dt
                        relative_time = 0.0
                    else:
                        delta = (timestamp_dt - base_time).total_seconds()
                        if delta < 0:
                            # Handle midnight crossover
                            delta += 24 * 3600
                        relative_time += delta
                        base_time = timestamp_dt
                except ValueError:
                    continue
            continue
        if current_section == 'top':
            # Handle timestamp line (if any). In provided sample, the top section starts with summary lines.
            # Therefore, we can assume that the relative_time increments by the top updates.

            # Detect the start of the process list by identifying the header line
            if re.match(r'^PID\s+USER\s+PR\s+NI\s+VIRT\s+RES\s+SHR\s+S\s+%CPU\s+%MEM\s+TIME\+\s+COMMAND$', line):
                in_process_list = True
                continue

            # Skip separator lines
            if re.match(r'^-+$', line):
                continue

            # If we're not yet in the process list, skip the line
            if not in_process_list:
                continue

            # Parse process information lines
            proc_info = re.split(r'\s+', line, maxsplit=10)
            if len(proc_info) < 11:
                continue
            try:
                pid = int(proc_info[0])
                user = proc_info[1]
                cpu_str = proc_info[8].replace(',', '.')  # Handle decimal point
                cpu = float(cpu_str) if re.match(r'^\d+(\.\d+)?$', cpu_str) else 0.0
                command = proc_info[10] if len(proc_info) > 10 else ""
            except (ValueError, IndexError):
                continue

            # Only record threads with SPIDs and CPU > 0
            if main_pid and pid in spids and cpu > 0.0:
                records.append({
                    'relative_time': relative_time,
                    'pid': pid,
                    'user': user,
                    'cpu': cpu,
                    'command': pid_to_cmd.get(pid, command)  # Use mapped command name if available
                })

    if not process_info_parsed:
        raise Exception("Process_Info section not found or improperly formatted in the topdump.txt file.")

    # Create DataFrame
    df = pd.DataFrame(records)
    print(f"Process name: {process_name}, Main PID: {main_pid}")
    return df, pid_to_cmd, process_name, main_pid

def plot_cpu_usage(df, process_identifier, process_name, pid_to_cmd, save_fig=False):
    """
    Plots CPU usage for the specified process and its threads on the same graph.

    Parameters:
    - df (pandas.DataFrame): DataFrame containing process and thread CPU usage data.
    - process_identifier (str or int): Process name or PID.
    - pid_to_cmd (dict): Dictionary mapping PID to command name.
    - save_fig (bool): Whether to save the plot as a PNG file.
    """

    plt.figure(figsize=(14, 8))

    # Ensure that each thread is only counted once per relative_time
    df_unique = df.drop_duplicates(subset=['relative_time', 'pid'])

    # Calculate total CPU usage by summing CPU of unique threads at each relative_time
    df_total = df_unique.groupby('relative_time')['cpu'].sum().reset_index()

    # Plot total CPU usage
    plt.plot(
        df_total['relative_time'],
        df_total['cpu'],
        label=f'Process {process_name}({process_identifier})',
        color='black',
        linewidth=2,
        linestyle='-'
    )

    # Plot each thread's CPU usage with command names
    for pid in df_unique['pid'].unique():
        df_pid = df_unique[df_unique['pid'] == pid]
        cmd_name = pid_to_cmd.get(pid, f'Thread {pid}')
        plt.plot(
            df_pid['relative_time'],
            df_pid['cpu'],
            label=f'{cmd_name} (PID {pid})',
            linewidth=1
        )

    plt.xlabel('Time (seconds)')
    plt.ylabel('CPU Usage (%)')
    plt.title(f'CPU Usage Over Time for Process {process_name} ({process_identifier}) and Its Threads')
    plt.legend(loc='upper left', bbox_to_anchor=(1, 1))
    plt.grid(True)
    plt.tight_layout()

    if save_fig:
        filename = f'{process_identifier}_cpu_usage.png'
        plt.savefig(filename)
        print(f"Plot saved as {filename}")

    plt.show()

def main():
    # Handle command-line arguments
    parser = argparse.ArgumentParser(description='Parse and plot CPU usage of the main process and its threads from topdump.txt.')
    parser.add_argument('-f', '--file', type=str, required=True, help='Path to the topdump.txt file')
    parser.add_argument('--save_fig', action='store_true', help='Save the generated plot as a PNG file')

    args = parser.parse_args()

    file_path = args.file
    save_fig = args.save_fig

    # Parse the topdump.txt file
    df, pid_to_cmd, process_name, main_pid = parse_topdump(file_path)

    # Plot CPU usage
    plot_cpu_usage(df, main_pid, process_name, pid_to_cmd, save_fig)

if __name__ == "__main__":
    main()
